import numpy as np
import cv2

from typing import List
from base_inference_engines import BaseInferenceEngine
from result import Result
from ais_bench.infer.interface import InferSession, MemorySummary
from utils import huawei_preprocess

confidence_threshold = 0.5
class HuaweiInferenceEngine(BaseInferenceEngine):
    def __init__(self):
        super().__init__()
        self.model = None

    def load_model(self, model_path: str, **kwargs) -> bool:
        # Implement Huawei model loading, e.g., using Ascend Toolkit
        self.model = InferSession(0, model_path)
        print("Loading model for Huawei...")
        return True

    def run_inference(self, img_path: str, **kwargs) -> List[Result]:
        # Implement Huawei-specific inference
        print("Running inference on Huawei...")
        input_image = cv2.imread(img_path) 
        result_list = []
        img_info = np.array([640,640,640,640], dtype=np.float16)
        input_image = huawei_preprocess(input_image)
        model_output = self.model.infer([input_image, img_info])
        box_out = model_output[0]
        box_out_num = model_output[1]
        idx = 0
        num_det = int(box_out_num[idx][0])
        boxout = box_out[idx][:num_det * 6].reshape(6, -1).transpose().astype(np.float32)  # 6xN -> Nx6
        for index, box in enumerate(boxout):
            if box[4] < confidence_threshold:
                continue
            bbox = Result.BBox(box[0], box[1], box[2], box[3])
            result = Result(bbox=bbox, confidence=box[4], label=box[5])
            result_list.append(result)
        # Format results to match the unified output structure
        return result_list